"""
VGG模型迁移学习案例

"""
import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
import tensorflow as tf
import numpy as np
from tensorflow.python import keras
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array, ImageDataGenerator
from tensorflow.python.keras.applications.vgg16 import VGG16, preprocess_input


class TransferModel:
    def __init__(self):
        # 创建训练数据和测试数据生成器对象
        self.train_generator = ImageDataGenerator(rescale=1.0 / 255.0)
        self.test_generator = ImageDataGenerator(rescale=1.0 / 255.0)

        # 数据相关
        self.train_path = "./data/train"
        self.test_path = "./data/test"
        self.target_size = (224, 224)
        self.batch_size = 32

        self.base_model = VGG16(include_top=False, weights='imagenet')

        self.label = {
            '0': 0,
            '1': 1,
            '2': 2,
            '3': 3,
            '4': 4
        }

    def load_data(self):
        """
        加载数据迭代器
        :return:
        """
        train_gen = self.train_generator.flow_from_directory(directory=self.train_path,
                                                             target_size=self.target_size,
                                                             class_mode="binary",
                                                             shuffle=True,
                                                             batch_size=self.batch_size)
        test_gen = self.test_generator.flow_from_directory(directory=self.test_path,
                                                           target_size=self.target_size,
                                                           class_mode="binary",
                                                           shuffle=True,
                                                           batch_size=self.batch_size)

        return train_gen, test_gen

    def refine_vgg_model(self):
        """
        调整VGG16模型
        :return:
        """
        x = self.base_model.output
        print(x)
        x = keras.layers.GlobalAveragePooling2D()(x)
        x = keras.layers.Dense(1024, activation=tf.nn.relu)(x)
        y_predict = keras.layers.Dense(5, activation=tf.nn.softmax)(x)
        model = keras.models.Model(inputs=self.base_model.input, outputs=y_predict)

        return model

    def freeze_model(self):
        """
        冻结模型
        :return:
        """
        for layer in self.base_model.layers:
            layer.trainable = False

    @staticmethod
    def compile(model):
        model.compile(optimizer=keras.optimizers.Adam(), loss=keras.losses.sparse_categorical_crossentropy,
                      metrics=["accuracy"])

    @staticmethod
    def fit_generator(model, train_gen, test_gen):
        check = keras.callbacks.ModelCheckpoint(filepath="./test_ckpt/transfer_{epoch:02d}-{acc:.2f}.h5",
                                                monitor="acc",
                                                save_best_only=True,
                                                save_weights_only=True,
                                                model="auto",
                                                period=1
                                                )

        model.fit_generator(generator=train_gen, epochs=2, validation_data=test_gen, callbacks=[check])

    def predict(self, model):
        model.load_weights("./test_ckpt/transfer_02-0.98.h5")
        image = load_img(path="./image/flower.jpg", target_size=self.target_size)
        image = img_to_array(image)
        image = image.reshape((1, image.shape[0], image.shape[1], image.shape[2],))
        image = preprocess_input(image)
        y_predictions = model.predict(image)
        print(self.label[str(np.argmax(y_predictions, axis=1)[0])])


if __name__ == '__main__':
    tm = TransferModel()

    # # 1. 获取数据
    # train_gen, test_gen = tm.load_data()

    # 2. 调整模型
    model = tm.refine_vgg_model()

    # # 3. 冻结模型
    # tm.freeze_model()
    #
    # # 4. 编译模型
    # tm.compile(model)
    #
    # # 5. 模型训练
    # tm.fit_generator(model, train_gen, test_gen)

    # 6. 模型预测
    tm.predict(model)
